use strict;
use warnings;

package Math::DirichletEstimator;

use base qw/Exporter/;

use Inline C => Config =>
    INC         => "-I/usr/local/lib/R/include/",
    LIBS        => '-lm -lRmath -llbfgs',
    CCFLAGS     => '-O3';

our @ISA = qw/Exporter/;

use Inline C => <<__C__;

#define MATHLIB_STANDALONE 1

#include <Rmath.h>
#include "lbfgs.h"

typedef struct {
    int         n, m;
    double**    xs;
} estimate_t;

int DIRICHLET() { return 1; }
int POLYA()     { return 2; }

void perl_array_to_c_double_array(AV* parr, double ** carr, int *carr_size) {
    
    int i, n;

    n = av_len(parr) + 1;
    
    /* allocate space for C array */
    *carr = (double*)malloc(sizeof(double) * n);
    *carr_size = n;

    /* fill up C array with values from Perl array */
    for (i = 0; i < n; i++) {
        (*carr)[i] = SvNV(*av_fetch(parr, i, 0));
    }

}

void perl_array_to_c_sv_array(AV* parr, SV *** carr, int *carr_size) {
    
    int i, n;

    n = av_len(parr) + 1;
    
    /* allocate space for C array */
    *carr = (SV**)malloc(sizeof(SV*) * n);
    *carr_size = n;

    /* fill up C array with values from Perl array */
    for (i = 0; i < n; i++) {
        (*carr)[i] = *av_fetch(parr, i, 0);
    }

}

double _log_dirichlet(double* x, double* alpha, int n, double sum_alpha) {

    int     i;
    double  r, log_x_i;

    r = 0;
    
    if (sum_alpha == -1) {
        sum_alpha = 0;
        for (i = 0; i < n; i++) {
            log_x_i = log(x[i]);
            r += alpha[i] * log_x_i - log_x_i - lgammafn(alpha[i]);
            sum_alpha += alpha[i];
        }
    }
    else {
        for (i = 0; i < n; i++) {
            log_x_i = log(x[i]);
            r += alpha[i] * log_x_i - log_x_i - lgammafn(alpha[i]);
        }
    }

    r += lgammafn(sum_alpha);

    return r;

}

double log_dirichlet(AV* av_x, AV* av_alpha) {

    int     i, n;
    double* alpha;
    double* x;
    double  r;

    n = av_len(av_x) + 1;

    alpha = (double*)malloc(sizeof(double) * n);
    x = (double*)malloc(sizeof(double) * n);
    for (i = 0; i < n; i++) {
        x[i] = SvNV(*av_fetch(av_x, i, 0));
        alpha[i] = SvNV(*av_fetch(av_alpha, i, 0));
    }

    r = _log_dirichlet(x, alpha, n, -1);

    free(alpha);
    free(x);

    return r;

}

double _log_polya(
    double* x, double* alpha, int n,
    double sum_alpha, double sum_lgamma_alpha) /* optional, -1 to ignore */
{

    int     i;
    double  r, sum_x;

    if (sum_alpha == -1 && sum_lgamma_alpha == -1) {
        sum_alpha = 0;
        sum_lgamma_alpha = 0;
        for (i = 0; i < n; i++) {
            sum_alpha += alpha[i];
            sum_lgamma_alpha += lgammafn(alpha[i]);
        }
    }
    else if (sum_alpha == -1) {
        sum_alpha = 0;
        for (i = 0; i < n; i++) sum_alpha += alpha[i];
    }
    else if (sum_lgamma_alpha == -1) {
        sum_lgamma_alpha = 0;
        for (i = 0; i < n; i++) sum_lgamma_alpha += lgammafn(alpha[i]);
    }

    sum_x = 0;
    for (i = 0; i < n; i++) sum_x += x[i];

    r = 0;
    r += lgammafn(sum_x + 1);
    r += lgammafn(sum_alpha);
    for (i = 0; i < n; i++) {
        r += lgammafn(alpha[i] + x[i]);
        r -= lgammafn(x[i] + 1);            
    }
    r -= lgammafn(sum_x + sum_alpha);
    r -= sum_lgamma_alpha;

    return r;

}

double log_polya(AV* av_x, AV* av_alpha) {

    int     i, n;
    double* alpha;
    double* x;
    double  r;

    n = av_len(av_x) + 1;

    alpha = (double*)malloc(sizeof(double) * n);
    x = (double*)malloc(sizeof(double) * n);
    for (i = 0; i < n; i++) {
        x[i] = SvNV(*av_fetch(av_x, i, 0));
        alpha[i] = SvNV(*av_fetch(av_alpha, i, 0));
    }

    r = _log_polya(x, alpha, n, -1, -1);

    free(alpha);
    free(x);

    return r;

}

double neg_log_dirichlet_likelihood(double** xs, int m, double* alpha, int n) {
    
    int     i;
    double  r, sum_alpha;

    r = 0;

    sum_alpha = 0;
    for (i = 0; i < n; i++) {
        sum_alpha += alpha[i];
    }
    
    for (i = 0; i < m; i++) {
        r -= _log_dirichlet(xs[i], alpha, n, sum_alpha);
    }

    return r;

}

void d_neg_log_dirichlet_likelihood(
    double** xs, int m, double* alpha, int n, double* d)
{

    int     i, j;
    double  sum_alpha, digamma_sum_alpha, digamma_alpha_i;

    for (i = 0; i < n; i++) {
        sum_alpha += alpha[i];
    }
    digamma_sum_alpha = digamma(sum_alpha);
    
    for (i = 0; i < n; i++) {
        d[i] = 0;
        digamma_alpha_i = digamma(alpha[i]);
        for (j = 0; j < m; j++) {
            d[i] -= log(xs[j][i]) - digamma_alpha_i + digamma_sum_alpha;
        }
    }

}

double neg_log_polya_likelihood(double** xs, int m, double* alpha, int n) {

    int     i;
    double  r, sum_alpha, sum_x, sum_lgamma_alpha;

    r = 0;

    sum_alpha = 0;
    sum_lgamma_alpha = 0;
    for (i = 0; i < n; i++) {
        sum_alpha += alpha[i];
        sum_lgamma_alpha += lgammafn(alpha[i]);
    }

    for (i = 0; i < m; i++) {
        r -= _log_polya(xs[i], alpha, n, sum_alpha, sum_lgamma_alpha);
    }

    return r;

}

void d_neg_log_polya_likelihood(
    double** xs, int m, double* alpha, int n, double* d)
{

    int     i, j;
    double  sum_alpha, sum_x, digamma_sum;

    sum_alpha = 0;
    for (i = 0; i < n; i++) {
        sum_alpha += alpha[i];
    }

    for (j = 0; j < m; j++) {
        sum_x = 0;
        for (i = 0; i < n; i++) {
            sum_x += xs[j][i];
        }
        digamma_sum = digamma(sum_alpha + sum_x);
        for (i = 0; i < n; i++) {
            d[i] = 0;
            d[i] -= digamma(sum_alpha);
            d[i] -= digamma(xs[j][i] + alpha[i]);
            d[i] += digamma_sum;
            d[i] += digamma(alpha[i]);
        }
    }

}

double dirichlet_eval_cb(
    void*           instance,
    const double*   alpha,
    double*         g,
    const int       n,
    const double    step)
{

    estimate_t* e = (estimate_t*)instance;
    
    d_neg_log_dirichlet_likelihood(e->xs, e->m, alpha, e->n, g);
    
    return neg_log_dirichlet_likelihood(e->xs, e->m, alpha, e->n);

}

double polya_eval_cb(
    void*           instance,
    const double*   alpha,
    double*         g,
    const int       n,
    const double    step)
{

    estimate_t* e = (estimate_t*)instance;

    fprintf(stderr, "polya_eval_cb called\\n");

    d_neg_log_polya_likelihood(e->xs, e->m, alpha, e->n, g);

    return neg_log_polya_likelihood(e->xs, e->m, alpha, e->n); 

}

int prgr_cb(
    void*                   instance,
    const lbfgsfloatval_t*  x,
    const lbfgsfloatval_t*  g,
    const lbfgsfloatval_t   fx,
    const lbfgsfloatval_t   xnorm,
    const lbfgsfloatval_t   gnorm,
    const lbfgsfloatval_t   step,
    int                     n,
    int                     k,
    int                     ls)
{

    fprintf(stderr, "%d: -log_lh = %g, |d_alpha| = %g, |alpha| = %g\\n",
        k, fx, gnorm, xnorm);
    fflush(stderr);
    return 0;

}

AV* estimate(AV* av_vector_refs, int model_type) {

    SV**                vector_refs;
    AV*                 av_vector;
    AV*                 av_alpha;
    int                 i, j, s, n_vectors, vector_size;
    double**            vectors;
    double*             alpha;
    lbfgs_parameter_t   lbfgs_param;
    estimate_t*         data;

    perl_array_to_c_sv_array(av_vector_refs, &vector_refs, &n_vectors);
    vectors = (double**)malloc(sizeof(double*) * n_vectors);

    for (i = 0; i < n_vectors; i++) {
        av_vector = (AV*)SvRV(vector_refs[i]);
        perl_array_to_c_double_array(av_vector, &vectors[i], &vector_size);
    }

    alpha = lbfgs_malloc(vector_size);
    for (i = 0; i < vector_size; i++) alpha[i] = 1;

    lbfgs_parameter_init(&lbfgs_param);
    lbfgs_param.epsilon = 1e-3;
    lbfgs_param.m = 10;

    data = (estimate_t*)malloc(sizeof(estimate_t));
    data->xs = vectors;
    data->n = vector_size;
    data->m = n_vectors;

    s = lbfgs(vector_size, alpha, NULL,
        model_type == DIRICHLET() ? &dirichlet_eval_cb : &polya_eval_cb,
        &prgr_cb, data, &lbfgs_param);
    
    av_alpha = newAV();
    av_extend(av_alpha, vector_size - 1);
    for (i = 0; i < vector_size; i++) av_store(av_alpha, i, newSVnv(alpha[i]));

    lbfgs_free(alpha);

    free(data);
    free(vector_refs);
    for (i = 0; i < n_vectors; i++) free(vectors[i]);
    free(vectors);

    return av_alpha;

}

__C__


our @EXPORT = qw/
    estimate
    log_dirichlet
    log_polya
    average_log_dirichlet
    average_log_polya
    DIRICHLET
    POLYA
/;

# give the average log dirichlet probability density on a set of x

sub average_log_dirichlet {

    my ($xs, $alpha) = @_;

    my $sum = 0;

    for my $x (@$xs) {
        $sum += log_dirichlet($x, $alpha);
    }

    return  $sum / @$xs;

}

sub average_log_polya {

    my ($xs, $alpha) = @_;

    my $sum = 0;

    for my $x (@$xs) {
        $sum += log_polya($x, $alpha);
    }

    return  $sum / @$xs;

}

1;

